from sklearn.cluster import KMeans
from sklearn.manifold import TSNE
from wordcloud import WordCloud
from excavator import LDAModeling
from gensim.models.ldamodel import LdaModel
from gensim.corpora import Dictionary
from typing import List
import pyLDAvis.gensim_models
import matplotlib.pyplot as plt
import numpy as np
import PIL.Image
import webbrowser
import os

# Enable Chinese-display
from pylab import mpl
mpl.rcParams['font.sans-serif'] = ['SimHei']

__paths = {
    'word_cloud':'./output_data/wordcloud/',
    'cluster':'./output_data/cluster/',
    'LDA':'./output_data/LDAPage/',
    'bar_chart':'./output_data/BarChart/'
}

for path in __paths.values():
    if not os.path.exists(path):
        os.mkdir(path)


def FormWordCloud(data, mask_path: str, font_path: str, title: str) -> PIL.Image:
    """
    @param:
        data: words separated with " " or word frequency dict
        mask_path: path of mask image(.jpg/.jpeg)
        font_path: path of font
        title: title of result
    @function:
        Form a wordcloud image(.jpg) with given text, mask and font
    @return:
        result image
    """
    if not os.path.exists(mask_path):
        raise FileNotFoundError('File "%s" not found.' % mask_path)
    if not os.path.exists(font_path):
        raise FileNotFoundError('File "%s" not found.' % font_path)

    mask = np.array(PIL.Image.open(mask_path))

    data_type = type(data)
    if data_type is str:
        res = WordCloud(mask=mask, font_path=font_path).generate(data)
    elif data_type is dict:
        res = WordCloud(mask=mask, font_path=font_path).generate_from_frequencies(data)
    else:
        raise ValueError('str or Dict expected, %s received.' % data_type)
    res_img = res.to_image()

    res_img.save(__paths['word_cloud']+ title + '.jpg')

    return res_img


def KMeansScatterGram(features_mat, n_clusters: int, title: str, show: bool = False) -> None:
    """
    @param:
        features_mat: matrix of feature values
        n_clusters: number of cluster to form
        title: title of result
    @function:
        Form a scatter gram(.png) using KMeans method
    @return:
        None
    """
    tsne = TSNE(n_components=2)
    decomposed = tsne.fit_transform(features_mat)
    km = KMeans(n_clusters=n_clusters, max_iter=5000)
    labels = km.fit_predict(decomposed)
    x = [i[0] for i in decomposed]
    y = [i[1] for i in decomposed]
    centers_x = [i[0] for i in km.cluster_centers_]
    centers_y = [i[1] for i in km.cluster_centers_]
    fig = plt.figure(figsize=(16, 16))
    plt.scatter(x, y, c=labels, marker='o')
    plt.scatter(centers_x, centers_y, c='red', marker='x')
    plt.xticks(())
    plt.yticks(())
    plt.savefig(__paths['cluster']+title+'.png')
    if show:
        plt.show()


def LDAVisualize(lda: LdaModel, corpus: List, dct: Dictionary, title: str, show: bool = False) -> None:
    vis_obj = pyLDAvis.gensim_models.prepare(lda, corpus, dct)
    pyLDAvis.save_html(vis_obj, __paths['LDA']+title+'.html')
    if show:
        webbrowser.open_new_tab(path)


def BarChart(data_x, labels_x, title: str, xlabel_name: str, ylabel_name: str, show: bool = False):
    plt.title(title)
    plt.xlabel(xlabel_name)
    plt.ylabel(ylabel_name)
    plt.bar(data_x, labels_x)
    plt.savefig(__paths['bar_chart']+title+'.jpg')
    if show:
        plt.show()
